{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a1a0cf63",
   "metadata": {},
   "source": [
    "# Plot t-SNE of speaker embeddings over SSL training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abdc0067",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Auto-reload imported modules from sslsv\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "# Load sslsv as a package from the parent folder\n",
    "import os\n",
    "import sys\n",
    "os.chdir('../..')\n",
    "sys.path.insert(1, os.path.join(sys.path[0], '../..'))\n",
    "\n",
    "# Embed fonts when saving figures as PDF\n",
    "import matplotlib\n",
    "matplotlib.rc('pdf', fonttype=42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7fccb6bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from notebooks.notebooks_utils import load_models, evaluate_models\n",
    "\n",
    "from sv_visualization import _filter_embeddings\n",
    "\n",
    "from sslsv.evaluations.CosineSVEvaluation import CosineSVEvaluation, CosineSVEvaluationTaskConfig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e2d2511",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy\n",
    "import pandas as pd\n",
    "\n",
    "from sklearn.manifold import TSNE\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.animation import FuncAnimation\n",
    "import seaborn as sns\n",
    "\n",
    "\n",
    "from glob import glob\n",
    "from tqdm import tqdm\n",
    "from pathlib import Path\n",
    "\n",
    "from IPython.display import Image, display"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d588741",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_tsne_embeddings(config, checkpoints, nb_speakers=10, nb_samples=150):\n",
    "    embeddings = []\n",
    "    labels = []\n",
    "\n",
    "    init = np.random.randn(nb_speakers*nb_samples, 2)\n",
    "\n",
    "    for checkpoint_path in tqdm(checkpoints, desc='Generating t-SNE of embeddings'):\n",
    "        checkpoint_name = Path(checkpoint_path).name\n",
    "\n",
    "        models = load_models(\n",
    "            [config],\n",
    "            checkpoint_name=checkpoint_name,\n",
    "        )\n",
    "        \n",
    "        evaluate_models(\n",
    "            models,\n",
    "            CosineSVEvaluation,\n",
    "            CosineSVEvaluationTaskConfig(__type__='sv_cosine'),\n",
    "            verbose=False\n",
    "        )\n",
    "\n",
    "        model = list(models.values())[0]\n",
    "\n",
    "        Z, y = _filter_embeddings(model.embeddings, nb_speakers=nb_speakers, nb_samples=nb_samples)\n",
    "        Z_2d = TSNE(\n",
    "            n_components=2,\n",
    "            init=init,\n",
    "            random_state=0\n",
    "        ).fit_transform(Z)\n",
    "\n",
    "        init = Z_2d\n",
    "\n",
    "        embeddings.append(Z_2d)\n",
    "        labels.append(y)\n",
    "\n",
    "    embeddings = np.array(embeddings).transpose(1, 0, 2) # (S, N, D) -> (N, S, D)\n",
    "    labels = np.array(labels)[0]\n",
    "    \n",
    "    return embeddings, labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18f64a0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_interpolated_embeddings(embeddings, nb_frames):\n",
    "    F = nb_frames\n",
    "    N, S, D = embeddings.shape\n",
    "\n",
    "    res = np.zeros((F, N, D))\n",
    "\n",
    "    for n in range(N):\n",
    "        fx = scipy.interpolate.interp1d(np.arange(0, S), embeddings[n, :, 0])\n",
    "        fy = scipy.interpolate.interp1d(np.arange(0, S), embeddings[n, :, 1])\n",
    "\n",
    "        for f, t in enumerate(np.linspace(0, S - 1, F)):\n",
    "            res[f, n, 0] = fx(t)\n",
    "            res[f, n, 1] = fy(t)\n",
    "            \n",
    "    return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "039f5da8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_animation(interpolated_embeddings, labels, frame_interval, nb_frames, checkpoints):\n",
    "    def update(i, ax):\n",
    "        ax.cla()\n",
    "        sns.scatterplot(\n",
    "            x=\"x\",\n",
    "            y=\"y\",\n",
    "            hue=\"Speaker\",\n",
    "            palette=sns.color_palette(\"hls\", len(np.unique(labels))),\n",
    "            data=df_list[i],\n",
    "            legend=\"full\",\n",
    "            alpha=0.6,\n",
    "            ax=ax\n",
    "        )\n",
    "        \n",
    "        ckpt_i = int(i // (nb_frames / len(checkpoints)))\n",
    "        step = int(checkpoints[ckpt_i].split('__step_')[1].split('.')[0])\n",
    "        ax.set_title(f'Training iteration: {step}', fontsize=12)\n",
    "        \n",
    "        ax.axis('off')\n",
    "        \n",
    "        legend = ax.legend(loc='lower left')\n",
    "        legend.set_title(\"Speaker\", prop={'size': 11})\n",
    "        legend.get_frame().set_linewidth(0)\n",
    "    \n",
    "    df_list = []\n",
    "    for Z in interpolated_embeddings:\n",
    "        df_list.append(pd.DataFrame({\n",
    "            \"Speaker\": labels,\n",
    "            \"x\": Z[:, 0],\n",
    "            \"y\": Z[:, 1],\n",
    "        }))\n",
    "    \n",
    "    dpi = 100\n",
    "    width = 1920 / dpi\n",
    "    height = 1080 / dpi\n",
    "    fig = plt.figure(figsize=(width, height))\n",
    "\n",
    "    ax = fig.gca()\n",
    "    anim = FuncAnimation(fig, update, frames=len(df_list), fargs=(ax,), interval=frame_interval)\n",
    "\n",
    "    plt.tight_layout()\n",
    "    padding = 0.01\n",
    "    fig.subplots_adjust(left=padding, bottom=padding, right=1-padding, top=0.95, wspace=None, hspace=None)\n",
    "    \n",
    "    plt.style.use(\"dark_background\")\n",
    "    return anim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dfb6375a",
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoints = glob('models/tests/simclr/model__step*')\n",
    "checkpoints = sorted(checkpoints, key=lambda f: int(f.split(\"__step_\")[1].split(\".\")[0]))\n",
    "\n",
    "config='./models/tests/simclr/config.yml'\n",
    "output=\"output.gif\"\n",
    "\n",
    "checkpoints_ = []\n",
    "checkpoints_ += checkpoints[  0:80:5]\n",
    "checkpoints_ += checkpoints[80:80*6:50]\n",
    "checkpoints_ += checkpoints[80*6::100]\n",
    "checkpoints_ += [checkpoints[-1]]\n",
    "checkpoints = checkpoints_\n",
    "\n",
    "nb_speakers = 10\n",
    "nb_samples = 150\n",
    "nb_frames = 30 * (len(checkpoints) - 1)\n",
    "frame_interval = 30\n",
    "\n",
    "print(\"Number of checkpoints:\", len(checkpoints))\n",
    "print(\"Number of frames:\", nb_frames)\n",
    "print(\"GIF duration:\", nb_frames*frame_interval/1000)\n",
    "print(\"FPS:\", 1000*1/frame_interval)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19d96d99",
   "metadata": {},
   "outputs": [],
   "source": [
    "# checkpoints = [\n",
    "#     'model__step_780.pt',\n",
    "#     'model__step_790.pt',\n",
    "# ]\n",
    "# nb_frames = 2\n",
    "# frame_interval = 1000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f27f09b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# embeddings, labels = generate_tsne_embeddings(\n",
    "#     config,\n",
    "#     checkpoints,\n",
    "#     nb_speakers,\n",
    "#     nb_samples\n",
    "# )\n",
    "\n",
    "# nb_frames = 1\n",
    "# interpolated_embeddings = generate_interpolated_embeddings(embeddings, nb_frames)\n",
    "# interpolated_embeddings = embeddings.transpose(1, 0, 2)\n",
    "\n",
    "anim = create_animation(\n",
    "    interpolated_embeddings,\n",
    "    labels,\n",
    "    frame_interval,\n",
    "    nb_frames,\n",
    "    checkpoints\n",
    ")\n",
    "anim.save(output, writer=\"pillow\")\n",
    "\n",
    "display(Image(output))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
